#include "crt.h"
#include "../dxut/dxstdafx.h"
#include <math.h>
#include <float.h>
#include <fstream>
#include <sys/timeb.h>
#include <time.h>
#include <fstream>


#define CRT_MAX_DEPTH	(3)		// (CRT_MAX_DEPTH - 1) extra bounces


using namespace std;


TSceneIntersection::TSceneIntersection()
{
	dist = FLT_MAX;
}


CRT::CRT()
{
	dirty = true;
	scene = NULL;

	srand((unsigned)time(NULL));
}

CRT::~CRT()
{
}



//
// ask the tracer to render one frame
// make sure the scene, the camera and the output settings are properly set
//
int CRT::Run()
{
	// if no input changes, then no output changes
	if (!dirty && output.buffer)
		return 1;

	// a lot of checking
	if (!scene)
	{
		MessageBox(0, L"No scene set to ray tracer!", L"Warning", MB_OK);
		return 0;
	}

	if (!scene->Check())
	{
		MessageBox(0, L"Scene error", L"Warning", MB_OK);
		return 0;
	}

	if (!camera.Check())
	{
		MessageBox(0, L"Camera error", L"Warning", MB_OK);
		return 0;
	}

	if (!output.Check())
	{
		MessageBox(0, L"Output error", L"Warning", MB_OK);
		return 0;
	}

	if (output.format != CRT_OF_BGR)
	{
		MessageBox(0, L"CRT supports CRT_OF_BGR format only for now!", L"Warning", MB_OK);
		return 0;
	}


	_timeb time;
	// time on starting point	
	_ftime( &time );
	output.stime_s	= time.time;
	output.stime_ms = time.millitm;

	// checking over, pre-run callback
	if (output.startfunc)
		output.startfunc();

	// buffer initialization
	int bufsize = 0;

	switch(output.format)
	{
	case CRT_OF_RGB:
	case CRT_OF_BGR:
		bufsize = output.imageh * output.imagew * 3;
		break;

	case CRT_OF_HDR_RGB:
	case CRT_OF_HDR_BGR:
		bufsize = output.imageh * output.imagew * 3 * 4;
		break;
	}

	// stat init
	output.nray			= 0;
	output.nrayaabb		= 0;
	output.nraytriangle = 0;
	scene->ResetStat();

	// start running the tracing code
	Trace();

	// post-run callback
	if (output.stopfunc)
		output.stopfunc();

	dirty = false;

	// time on ending point
	_ftime( &time );
	output.timeused = float(time.time - output.stime_s) + 0.001f * (time.millitm - output.stime_ms);

	return 1;
}


#define POSITION_SHIFT	(0.0001f)

// phong model. 
VEC3 CRT::Shade(RAY& ray, VEC3& pos, VEC3& normal, TCRTMaterial& material, int depth)
{
	VEC3 hdrColor(0);

	// eye direction
	VEC3 E = -ray.direction;

	// normal
	VEC3 N = normal;

	//light number
	int lightNum;

	vector<VEC3> rays;
	// ambient, diffuse and specular
	for (vector<TCRTLight>::iterator it = scene->lights.begin(); it != scene->lights.end(); it++)
	{
		// TONGJI_TODO


		if(it->type == CRT_LT_SKY){
			int div = 16;
			double u = 0;
			double v = -M_PI;
			double u_delta = M_PI/((double)div*2);
			double v_delta = 2*M_PI/((double)div);

			double rad = 1.0;
			VEC3 top = VEC3(0.0,rad,0.0);


			int i;
			for ( i=0; i<div; i++ )
			{
				v = -M_PI;
				for ( int j=0; j<div; j++ )
				{
					v += v_delta;
					if(rand()%10==0)
					{
						double x = rad * cos(u) * cos(v);
						double y = rad * sin(u);
						double z = rad * cos(u) * sin(v) ;
						VEC3 pnt( x, y, z );
						rays.push_back(pnt);
					}
				}
				u += u_delta;
			}
			rays.push_back(top);
		}
		else if(it->type == CRT_LT_DIRECTIONAL)
		{
			rays.push_back(pos);
		}

		lightNum = rays.size(); 
		for(int i = 0;i < rays.size();++i)
		{

			// light direction 
			VEC3 L = it->GetDirection(pos,rays[i]);

			//ofstream temp("output.txt");

			//double step0 = clock();
			// light direction mirror according to Normal
			VEC3 M = L.Mirror(N);
			//double step1 = clock();

			//temp<<"Mirror "<<(step1-step0)*1000000000000<<endl;

			hdrColor += material.ambient.GetVec3();				// ambient


			float atten_shadow = 1.0f;	// attenuation due to shadow ray blocking

			if (output.shadow)
			{
				// shadow ray testing
				RAY shadowRay;
				shadowRay.origin	= pos;
				shadowRay.direction	= L;
				shadowRay.origin += shadowRay.direction * POSITION_SHIFT;	// a little bit shift

				vector<TSceneIntersection> infos;

				// objects are treated as with full opacity
				if (scene->TraceRay(shadowRay, infos, kStopOnFirstHit))
					atten_shadow = 0.0f;
				else
					atten_shadow = 1.0f;
			}
			//double step2 = clock();
			//temp<<"traceray "<<(step2-step1)*1000000000000<<endl;

			if ( atten_shadow >= CRT_MATH_EPSILON )
			{		
				VEC3 dif_spe(0);

				dif_spe += material.diffuse.GetVec3() * N.Dot_Saturate(L);	// diffuse
				float dot = E.Dot_Saturate(M);
				dif_spe += material.specular.GetVec3() * powf(dot, material.shininess);	// specular

				hdrColor += dif_spe * atten_shadow;
			}

			hdrColor *= it->color;	// light color
			// double step3 = clock();
			//temp<<"matrices "<<(step3-step2)*1000000000000<<endl;
		}

	}

	//reduce light intensity

	hdrColor.x/=lightNum;
	hdrColor.y/=lightNum;
	hdrColor.z/=lightNum;

	return hdrColor;


}	

bool CRT::TraceRay(RAY& ray, int depth, VEC3& shading)
{
	shading = VEC3(0);

	vector<TSceneIntersection> infos;
	bool hit = scene->TraceRay(ray, infos, kNearestHit);

	if (hit)
	{
		TCRTMesh& mesh = *(scene->meshes[infos[0].meshIndex]);

		TCRTVertexFull& p0 = mesh.vertices_f[infos[0].index[0]];
		TCRTVertexFull& p1 = mesh.vertices_f[infos[0].index[1]];
		TCRTVertexFull& p2 = mesh.vertices_f[infos[0].index[2]];

		// position
		VEC3 pos = p0.position + (p1.position - p0.position) * infos[0].u + (p2.position - p0.position) * infos[0].v;
		// normal
		VEC3 normal = p0.normal + (p1.normal - p0.normal) * infos[0].u + (p2.normal - p0.normal) * infos[0].v;						
		normal.Normalize();

		shading = Shade( ray, pos, normal, mesh.material, depth );
	}

	return hit;
}

// clusterWidth/clusterHeight usually equal to mClusterHeight/mClusterWidth, but the last column/row may not. 
void CRT::TraceCluster(int clusterX, int clusterY, int clusterWidth, int clusterHeight)
{
	// outputing - direction of each ray
	VEC3 pixelPosition = camera.position;
	VEC3 pixelDirection;

	D3DXVECTOR3 src, dest;
	D3DXMATRIX worldMatrix;
	D3DXMatrixIdentity(&worldMatrix);

	unsigned char* pScanlineStarter = (unsigned char*)output.buffer + clusterY * output.imagew * 3 + clusterX * 3;
	for (int row = clusterY; row < clusterY + clusterHeight; row++)
	{
		unsigned char* pCurrentPixel = pScanlineStarter;
		for (int col = clusterX; col < clusterX + clusterWidth; col++)
		{
			src = D3DXVECTOR3( col + 0.5f, output.imageh - 1 - row + 0.5f, 0 );
			D3DXVec3Unproject( &dest, &src, &mViewport, &mProjectionMatrix, &mViewMatrix, &worldMatrix );

			pixelDirection = VEC3(dest.x, dest.y, dest.z) - pixelPosition;
			pixelDirection.Normalize();

			RAY ray; 
			ray.origin = pixelPosition;
			ray.direction = pixelDirection;
			ray.inObj = false;

			VEC3 shading;
			TraceRay( ray, 0, shading );
			COLOR color = COLOR( shading, true );

			*pCurrentPixel++ = color.b;
			*pCurrentPixel++ = color.g;
			*pCurrentPixel++ = color.r;
		}

		// Move to next row
		pScanlineStarter += output.imagew * 3;
	}
}


void CRT::ProgressUpdate(
	int nDirtyRectX, 
	int nDirtyRectY, 
	int nDirtyRectWidth,
	int nDirtyRectHeight,
	int nPrimaryRay, 
	float fFinishedPercent)
{
	// time update
	_timeb time;
	_ftime( &time );
	output.timeused = float(time.time - output.stime_s) + 0.001f * (time.millitm - output.stime_ms);

	// stat update
	output.nray			= nPrimaryRay;
	output.nrayaabb		= scene->GetRayAABBTestNum();
	output.nraytriangle = scene->GetRayTriangleTestNum();
	output.raytritime	= scene->GetRayTriangleTestTime();
	output.nDirtyRectX	= nDirtyRectX;
	output.nDirtyRectY	= nDirtyRectY;
	output.nDirtyRectWidth	= nDirtyRectWidth;
	output.nDirtyRectHeight	= nDirtyRectHeight;


	// progress callback
	if (output.progressfunc)
	{
		output.progressfunc(fFinishedPercent, 0);
	}
}

// Get view and projection matrix
void CRT::UpdateViewProjectionMatrix()
{
	mViewport.X = 0;
	mViewport.Y = 0;
	mViewport.Width = output.imagew;
	mViewport.Height = output.imageh;
	mViewport.MinZ = 0.0f;
	mViewport.MaxZ = 1.0f;

	for (int r = 0; r < 4; r++)
	{
		for (int c = 0; c < 4; c++)
		{
			mProjectionMatrix(r, c) = camera.projM.m[r * 4 + c];
			mViewMatrix(r, c) = camera.viewM.m[r * 4 + c];
		}
	}
}


void CRT::Trace()
{
	UpdateViewProjectionMatrix();

	int clusterRow = (output.imageh + mClusterHeight - 1) / mClusterHeight;
	int clusterCol = (output.imagew + mClusterWidth - 1) / mClusterWidth;
	int totalCluster = clusterRow * clusterCol;

	int nUpdateEveryCluster = max(1, totalCluster / 100);

	int clusterX = 0;
	int clusterY = 0;
	int clusterWidth = 0;
	int clusterHeight = 0;

	int nPrimaryRay = 0;
	int nFinishedCluster = 0;

	int remainedHeight = output.imageh;
	for (int row = 0; row < clusterRow; row++)
	{
		clusterY = output.imageh - remainedHeight;
		clusterHeight = min(CRT::mClusterHeight, remainedHeight);

		int remainedWidth = output.imagew;
		for (int col = 0; col < clusterCol; col++)
		{
			clusterX = output.imagew - remainedWidth;
			clusterWidth = min(CRT::mClusterWidth, remainedWidth);

			TraceCluster(clusterX, clusterY, clusterWidth, clusterHeight);
			remainedWidth -= clusterWidth;

			// Stat
			nFinishedCluster++;
			nPrimaryRay += clusterWidth * clusterHeight;

			// Progress update
			if (0 == nFinishedCluster % nUpdateEveryCluster)
			{
				ProgressUpdate(clusterX, clusterY, clusterWidth, clusterHeight, nPrimaryRay, 1.0f * nFinishedCluster / totalCluster);
			}
		}

		remainedHeight -= clusterHeight;
	}
}